from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_zhipu import ChatZhipuAI
from component.vector_databases import VectorDB
from component.embedding import ZhipuEmbedding
from langchain.prompts import  PromptTemplate
from langchain_core.output_parsers import StrOutputParser

import os


class Qwen2LLM():
    def __init__(self, model_name="gpt-3.5.turbo", temperature:float=0.9) -> None:
        
        # 初始化大模型
        self.model_name = model_name
        self.temperature = temperature
        llm = ChatOpenAI(base_url=os.getenv("QWEN2_BASE_URL"),
                 api_key=os.getenv("QWEN2_API_KEY"),
                 temperature=0.7)
        self.model = llm
        # 加载向量数据库
        self.db = VectorDB()
        self.db.load_vector(path="db")
        self.embedding_model = ZhipuEmbedding()

    def chat(self, text:str):
        #这里利用输入的问题与向量数据库里的相似度来匹配最相关的信息，填充到输入的提示词中
        template="""
        使用以上下文来回答用户的问题。如果你不知道答案，就说你不知道。总是使用中文回答。
        问题: {text}
        可参考的上下文：
        ···
        {info}
        ···
        如果给定的上下文无法让你做出回答，请回答数据库中没有这个内容，你不知道。
        有用的回答:"""
        info=self.db.query(text,self.embedding_model,1)
        prompt=PromptTemplate(template=template,input_variables=["text","info"]).format(text=text,info=info)
        # print(prompt)
        res=self.model.invoke(prompt)
        return  res.content


class OpenAILLM():
    def __init__(self, model_name="gpt-3.5.turbo", temperature:float=0.9) -> None:
        
        # 初始化大模型
        self.model_name = model_name
        self.temperature = temperature
        llm = ChatOpenAI(base_url=os.getenv("OPENAI_BASE_URL"),
                 api_key=os.getenv("OPENAI_API_KEY"),
                 temperature=0.7)
        self.model = llm
        # 加载向量数据库
        self.db = VectorDB()
        self.db.load_vector(path="db")
        self.embedding_model = ZhipuEmbedding()

    def chat(self, text:str):
        #这里利用输入的问题与向量数据库里的相似度来匹配最相关的信息，填充到输入的提示词中
        template="""使用以上下文来回答用户的问题。如果你不知道答案，就说你不知道。总是使用中文回答。
        问题: {text}
        可参考的上下文：
        ···
        {info}
        ···
        如果给定的上下文无法让你做出回答，请回答数据库中没有这个内容，你不知道。
        有用的回答:"""
        info=self.db.query(text,self.embedding_model, 1)
        prompt=PromptTemplate(template=template,input_variables=["text","info"]).format(text=text,info=info)
        res=self.model.invoke(prompt)
        return  res.content


class ZhipuLLM():
    def __init__(self, model_name="glm-4-0520", temperature:float=0.7):
        self.model_name = model_name
        self.temperature = temperature
        llm = ChatZhipuAI(
            model=model_name,
            temperature=temperature,
            api_key=os.getenv("ZHIPUAI_API_KEY")
        )
        self.model = llm
        # 加载向量数据库
        self.db = VectorDB()
        self.db.load_vector(path="db")
        self.embedding_model = ZhipuEmbedding()

        # 输出
        self.output_parser = StrOutputParser()

        # 上下文对话

        
    
    def chat(self, text:str):
        """
        聊天
        """

        template = """
        仅根据所提供的上下文回答以下问题:
        <context>
        {context}
        </context>

        问题：{input}
        """
        prompt = ChatPromptTemplate.from_template(template)
        info = self.db.query(text, self.embedding_model, 1)
        llm = prompt | self.model | self.output_parser
        res_text = llm.invoke({"context": info, "input": text})
        return res_text
    
    def chat_for_template(self, text:str, prompt_template=None):
        """根据模板进行聊天"""
        pass

    def chat_robot(self, text:str):
        pass
